
(* ===================================================== *)
(* ============== CS4212 Compiler Design ============== *)
(*     TypeChecking of Jlite programs *)
(*     A Simplified version  *)
(*     Submitting at your own risk  *)
(* ===================================================== *)

open Jlite_structs
open List

(* Compare two variable ids *) 	
let compare_var_ids v1 v2 =
	match v1,v2 with
	| SimpleVarId id1, SimpleVarId id2 -> 
		((String.compare id1 id2) == 0)
	| SimpleVarId id1, TypedVarId (id2,t,s) -> 
		((String.compare id1 id2) == 0)	
	| TypedVarId (id1,t,s), SimpleVarId id2 -> 
		((String.compare id1 id2) == 0)		
	| TypedVarId (id1,t1,s1), TypedVarId (id2,t2 ,s2) ->
		((String.compare id1 id2) == 0) && (s1 == s2)

(* Compare two variable decl lists *)
(* Used by checking method overloading *)
let rec compareVarDecls v1 v2 = match v1, v2 with
| [], [] -> true
| [], _
| _, [] -> false
| x::xs, y::ys -> x = y && compareVarDecls xs ys

(* Given a md_decl list and a method, returns true if the method with the same name and parameters exist in the md list*)
(* Otherwise returns false *)
(* Used by checking method overloading *)
let rec listContainsMethod (mdDecls: md_decl list) (mdDecl:md_decl): bool =
	if List.length mdDecls > 0 then
		let aMdDecl = List.hd mdDecls in
			let aMdName = aMdDecl.jliteid in
				let aMdParams = aMdDecl.params in
					if (aMdName = mdDecl.jliteid) && (compareVarDecls aMdParams (mdDecl.params)) then
						true
					else
						listContainsMethod (List.tl mdDecls) mdDecl
	else false
(* Find the declared type of a variable *) 		
let rec find_var_decl_type 
	(vlst: var_decl list) (vid:var_id) =
  match vlst with
    | [] -> (Unknown, SimpleVarId "") 
    | (t,v)::lst -> 
		if (compare_var_ids v vid) 
		then (t,v) 
		else (find_var_decl_type lst vid)

(* Check if a variable id exists *)
let exists_var_id 
	(vlst: var_decl list) (vid: var_id) : bool =
	let is_eq_id ((t,v): var_decl):bool =
		(compare_var_ids v vid) 
	in (List.exists is_eq_id vlst) 	

(* Check if the declaration of a class exists *) 			  
let exists_class_decl 
	((cm,clst): jlite_program) (cid:class_name) =
	(List.exists 
		(fun (cname,cvars,cmtd) ->
			(String.compare cname cid) == 0)
	clst)

(* HELPER FUNCTION: method exp *)
let rec argList_contains_unknown_arg (argList:(jlite_type * jlite_exp) list) : bool = 
	match argList with
	| anArg::restArgs ->
		if (fst anArg) = Unknown then
			true
		else
			argList_contains_unknown_arg restArgs
	| [] -> false

(**********************************************************)
(***********START: HELPER FUNCTION: method exp*************)
(**********************************************************)
(* Find the class with classname in the program *) 			  
let find_class_decl_in_program
	((cm,clst): jlite_program) (cid:class_name) =
	(List.find 
		(fun (cname,cvars,cmtd) ->
			(String.compare cname cid) == 0)
	clst) 

let rec get_md_class (p:jlite_program) (env: var_decl list) (currClassName:class_name) (mdStr:jlite_exp): class_decl =
		match mdStr with
		| FieldAccess (exp1,vid1) -> get_md_class p env currClassName exp1
		| ObjectCreate (cname1) -> find_class_decl_in_program p cname1
		| MdCall (exp1, exp1List) -> get_md_class p env currClassName exp1
		| Var varId ->
			let (varType, var) = find_var_decl_type env varId in
				begin
					match varType with
					| ObjectT ot -> (find_class_decl_in_program p ot)
					| _ -> find_class_decl_in_program p currClassName
				end
		| _ -> raise Not_found

let rec get_md_id (mdStr:jlite_exp): var_id = 
	match mdStr with
	| FieldAccess (exp1, vid1) -> vid1
	| Var varId -> varId
	| _ -> SimpleVarId ""

let rec compare_var_decl_list_and_typed_exp_list (varDeclList:var_decl list) (typedArgList: (jlite_type*jlite_exp) list):bool = 
	match varDeclList, typedArgList with
	| [],[] -> true
	| [], _
	| _, [] -> false
	| x::xs, (yType, yExp)::ys -> fst x = yType && compare_var_decl_list_and_typed_exp_list xs ys

let rec find_method_return_type (mdList:md_decl list) (mdid:var_id) (typedArgList:(jlite_type * jlite_exp) list): jlite_type = 
  match mdList with
  	| [] -> Unknown
  	| aMethod::remainingMtds ->
  		if ((aMethod.jliteid) = mdid && (compare_var_decl_list_and_typed_exp_list aMethod.params typedArgList)) then
  			aMethod.rettype
  		else find_method_return_type remainingMtds mdid typedArgList

let rec get_md_decl_given_md_str (p:jlite_program) (env: var_decl list) 
	(currClassName:class_name) (mdStr:jlite_exp) (typedArgList:(jlite_type*jlite_exp) list): jlite_type =
	let (cname, varList, mdList) = get_md_class p env currClassName mdStr in
		let mdid = get_md_id mdStr in
			if mdid = SimpleVarId "" then
				Unknown
			else
  			find_method_return_type mdList mdid typedArgList

(**********************************************************)
(***********END: HELPER FUNCTION: method exp*************)
(**********************************************************)




(* Annotate a list of variable declarations with their scope *)	
let rec create_scoped_var_decls
	(vlst: var_decl list) (scope:int) =
	let helper ((vt,vid):var_decl) =
		match vid with
		| SimpleVarId id -> 
			(vt, TypedVarId (id, vt, scope))
		| TypedVarId (id,t,s) -> 
			(vt,TypedVarId (id, vt, scope))
	in (List.map helper vlst)
  
(* Type check a list of variable declarations 
  1) Determine if all object types exist
  2) Find and return duplicate variable names -- DONE

    This type check function is a simple version in 
    which there is only one variable declaration in 
    the declaration list
	The function returns the 
	 typecheck result (true or false) and an error message.
*)  
let rec type_check_var_decl_list
	(p: jlite_program) 
	(vlst: var_decl list) =
	let rec helper 
		(vlst: var_decl list) :jlite_type list =
		if List.length vlst > 0 then
			match List.hd vlst with
			| (typ,vid) -> 
				begin
  			match typ with
  			| ObjectT cname -> 
  	        (* check if the declared class name exists *)
  				if (exists_class_decl p cname) 
  					then helper (List.tl vlst)
  					(* return the undefined type *)
  					else typ::(helper (List.tl vlst))
  			(* Primitive type *)
  			| StringT -> (helper (List.tl vlst))
  			| BoolT -> (helper (List.tl vlst))
  			| IntT -> (helper (List.tl vlst))
  			| VoidT -> (helper (List.tl vlst))
  			| typ -> ( typ::helper (List.tl vlst )) 
  			end
		else []
	in match ( helper vlst) with
		| [] ->  (true,"")
		| lst -> (false, ("Undefined types: " 
				^ (string_of_list lst string_of_jlite_type ",")))

 
(* Type check a list of method declarations 
  1) Determine if there is illegal overloading
  2) Find and return overloaded method names	

    This simplified version always returns true
	---  DONE  ---
*)
let rec type_check_md_overloading 
	(classid: class_name) (mdlst: md_decl list): (bool* string) =
		let rec helper (existMdList:md_decl list) (todoMdList: md_decl list) : var_id list=
			match todoMdList with
			| (aMd:md_decl)::(remainingMds:md_decl list) -> 
				if (listContainsMethod existMdList aMd) then
					aMd.jliteid :: (helper (aMd::existMdList) remainingMds)
				else
					helper (aMd::existMdList) remainingMds
			| [] -> []
		in match (helper [] mdlst) with
		| [] ->  (true,"")
		| lst -> (false, (string_of_list lst string_of_var_id ", "))

(* Type check an expression *)
(* Return the type of the Expression 
    and a new TypedExpession *)  
let rec type_check_expr 
	(p: jlite_program) (env: var_decl list) 
	(classid: class_name) (exp:jlite_exp) = 
	let rec helper e 
	: (jlite_type * jlite_exp) =
		match e with
		| BoolLiteral v -> (BoolT, e)
		| IntLiteral v -> (IntT, e)
		| StringLiteral v -> (StringT, e)
		| ThisWord -> 
			((ObjectT classid), TypedExp (e,(ObjectT classid)))
		| NullWord -> 
			((ObjectT "null") , TypedExp (e,(ObjectT "null")))
		| Var v -> 
			let (vtyp,vid) =(find_var_decl_type env v) in
			(vtyp, TypedExp (Var vid,vtyp)) 
		| ObjectCreate c -> 
			if (exists_class_decl p c) 
			then (ObjectT c, TypedExp (e, (ObjectT c)))
			else (Unknown, e)
		 (* Handle other expresion types ---- TODO ---- *)
		| UnaryExp (op, ue) -> 
			let eType = helper ue in
				if (fst eType)=IntT then
					eType
				else
					(Unknown, e)
		| BinaryExp (op, exp1, exp2) -> 
			begin
			match op with
			| BooleanOp bOp
			| RelationalOp bOp->
				if (fst (helper exp1)<>Unknown) && (fst (helper exp2)<>Unknown) then
					(BoolT, TypedExp (e, BoolT))
				else
					(Unknown, e)
			| AritmeticOp aOp->
				if (fst (helper exp1)=IntT) && (fst (helper exp2)=IntT) then
					(IntT, TypedExp (e, IntT))
				else (Unknown, e)
			| _ -> (Unknown, e)
			end
		| FieldAccess (obj,field) -> 
			if (exists_var_id env field) then
				let objType = fst (helper obj) in
  				let fieldType = fst (helper (Var field)) in
  					match objType with
  					| ObjectT ot ->
    					if (fieldType <> Unknown)
    					then (fieldType, TypedExp (e, fieldType))
    					else (Unknown, e)
  					| _ -> (Unknown, e)
			else (Unknown, e)
		| MdCall (mdStr, varList) ->  (*TODO: write mdcall typecheck....*)
				(*Firstly, type check the method argument list *)
				let newVarList = List.map helper varList in
					(* iterate var list; if it contains unknown element then return unknown as the type of this method *)
					if (argList_contains_unknown_arg newVarList) = true then
						(Unknown, e)
					else (* the argument list is clean, proceed to varify with the declared type of the method *)
					let mdType = try (get_md_decl_given_md_str p env classid mdStr newVarList) with Not_found -> Unknown in
  					(mdType, e)
		| _ -> (Unknown, e) 
	  in  helper exp

(* HELPER: after type check a stmt list, check if stmt list contains return statement *)
(* If statement contains return stmt, return the return Some <stmt type> *)
(* else, return None *)
let rec check_for_return_type (stmtList:(jlite_type option * jlite_stmt) list) = 
	match stmtList with 
	| [] -> None
	| (option, stmtExpr)::restStmt ->
		if option <> None then
			option
		else check_for_return_type restStmt

(* Type check a list of statements and determine the return type.
   Exceptions are thrown when a statement does not type check 
   and when dead code is found
*)  
let rec type_check_stmts 
	(p: jlite_program)(env: var_decl list) 
	(classid: class_name) 
	(mthd: md_decl) 
	(stmtlst:jlite_stmt list)
	(rettype: jlite_type option) 
	: (jlite_type option *(jlite_stmt list))  =
	match stmtlst with
	| [] -> (rettype,[])
	| s::tail_lst -> 
		let rec helper s 
		: (jlite_type option * jlite_stmt) =
		match s with
		| ReturnStmt e ->  
			let (expr_type,exprnew) = 
			 (type_check_expr p env classid e) in
			begin
			match expr_type with
			| Unknown -> 
				failwith 
				("\nType-check error in " 
				^ classid ^ "." ^ string_of_var_id mthd.jliteid 
				^ ". Return expression fails:\n" 
				^ string_of_jlite_stmt s ^ "\n")
			| _ ->  (Some expr_type, ReturnStmt exprnew)
			end
		| ReturnVoidStmt ->  
			(Some VoidT,ReturnVoidStmt)
		| ReadStmt id -> 
			let (idtype,scopedid) = (find_var_decl_type env id) in
			begin
			match idtype with
			| ObjectT _ | Unknown  -> 
				failwith 
				("\nType-check error in " 
				^ classid ^ "." ^ string_of_var_id mthd.jliteid 
				^ ". Read statement fails:\n" 
				^ string_of_jlite_stmt s ^ "\n")
			| _ ->  (None,ReadStmt scopedid)
			end
		| PrintStmt e -> 
			let (expr_type,exprnew) = 
			 (type_check_expr p env classid e) in
			begin
			match expr_type with
			| Unknown | ObjectT _ -> 
				failwith 
				("\nType-check error in " 
				^ classid ^ "." ^ string_of_var_id mthd.jliteid 
				^ ". Statement fails:\n" 
				^ string_of_jlite_stmt s ^ "\n")
			| _ ->  (None,PrintStmt exprnew)
			end
		(* _ -> Handle other Statement types
		  ---- TODO ---- *)
		| IfStmt (e, stmtList1, stmtList2) ->
			let (expr_type, exprnew) = (type_check_expr p env classid e) in
				begin
				match expr_type with
				| BoolT -> 
					let retType1 = check_for_return_type (List.map helper stmtList1) in
						let retType2 = check_for_return_type (List.map helper stmtList2) in
							if retType1 = retType2 then
								(retType1, IfStmt (exprnew, stmtList1, stmtList2))
							else
								(None, IfStmt (exprnew, stmtList1, stmtList2))
				| _ -> 
					failwith 
					("\nType-check error in " 
					^ classid ^ "." ^ string_of_var_id mthd.jliteid 
					^ ". If statement fails:\n" 
					^ string_of_jlite_stmt s ^ "\n")
				end
		| WhileStmt (e, stmtList) ->
			let (expr_type, exprnew) = (type_check_expr p env classid e) in
				begin
				match expr_type with
				| BoolT -> 
					let retType = check_for_return_type (List.map helper stmtList) in
						(retType, WhileStmt (exprnew, stmtList))
				| _ -> 
					failwith 
					("\nType-check error in " 
					^ classid ^ "." ^ string_of_var_id mthd.jliteid 
					^ ". While statement fails:\n" 
					^ string_of_jlite_stmt s ^ "\n")
				end
		| AssignStmt (vid, exp) ->
			let (idType,scopedId) = (find_var_decl_type env vid) in
				let (expType, typedExp) = type_check_expr p env classid exp in
					if idType = expType then
						(None, AssignStmt (scopedId, typedExp))
					else
						failwith 
  					("\nType-check error in " 
  					^ classid ^ "." ^ string_of_var_id mthd.jliteid 
  					^ ". AssignStmt fails:\n" 
  					^ string_of_jlite_stmt s ^ "\n")
		| AssignFieldStmt (fieldExp, assignExp) ->
			let (fieldExpType, fieldExpNew) = type_check_expr p env classid fieldExp in
				let (assignExpType, assignExpNew) = type_check_expr p env classid assignExp in
					if fieldExpType = assignExpType then
						(None, AssignFieldStmt (fieldExpNew, assignExpNew))
					else
						failwith 
  					("\nType-check error in " 
  					^ classid ^ "." ^ string_of_var_id mthd.jliteid 
  					^ ". AssignFieldStmt fails:\n" 
  					^ string_of_jlite_stmt s ^ "\n")
		| MdCallStmt mdExp -> 
			let (mdExpType, newMdExp) = (type_check_expr p env classid mdExp) in
				if mdExpType = Unknown then
					failwith 
  					("\nType-check error in " 
  					^ classid ^ "." ^ string_of_var_id mthd.jliteid 
  					^ ". MdCallStmt fails:\n" 
  					^ string_of_jlite_stmt s ^ "\n")
				else
					(None, MdCallStmt newMdExp)
	  in let (newrettype,newstmt) = ( helper s) in
	  match newrettype,tail_lst with
		| Some t, head::tail -> 
			failwith 
			("\nType-check error in " ^ classid ^ "." 
			 ^ string_of_var_id mthd.jliteid 
			 ^ ". Dead Code:\n" 
			 ^ (string_of_list tail_lst string_of_jlite_stmt "\n" ^ "\n")) 
		| _,_ ->  
			let (rettype,stmts) = 
				(type_check_stmts p env classid mthd tail_lst newrettype) in
				(rettype,(newstmt::stmts))
  
(* TypeCheck a JLite Method Declaration *)
let type_check_mthd_decl p env cname m : md_decl = 
	let mthdenv = 
		List.append m.params m.localvars in 
	let (retval, errmsg) = 
		(type_check_var_decl_list p mthdenv)
	in if (retval == false) 
		then failwith 
		 ("\nType-check error in " ^ cname ^ "." 
		  ^ string_of_var_id m.jliteid 
		  ^ " parameter or local variables declarations.\n"
		  ^ errmsg ^ "\n")
		else
		let scopedEnv = List.append 
				(create_scoped_var_decls mthdenv 2) env in 
		(* TypeCheck the body of the method *)
			let (rettyp,newstmts) = 
				(type_check_stmts p scopedEnv cname m m.stmts None) in
		(* TypeCheck the return type of the method *)
			let _ = match rettyp,m.rettype with
			| None, VoidT -> true
			| Some VoidT, VoidT -> true
			| None, t -> 
				failwith 
				("\nType-check error in " ^ cname ^ "." 
				^ string_of_var_id m.jliteid 
				^ ". This method must return a result of type "
				^ string_of_jlite_type m.rettype ^ ". \n")
			| Some (ObjectT t1), (ObjectT t2) -> 
				if ((String.compare t1 t2) != 0) 
				then failwith 
					("\nType-check error in " ^ cname ^ "." 
					^ string_of_var_id m.jliteid 
					^ ". Type mismatch. Return type of method " 
					^ "is different from declared type "
					^ string_of_jlite_type m.rettype ^ t1 ^ ". \n")
				else true
			| Some t1, t2 -> 
				if (t1!= t2) 
				then failwith 
					("\nType-check error in " ^ cname ^ "." 
					^ string_of_var_id m.jliteid 
					^ ". Type mismatch. Return type of method "
					^ "is different from declared type "
					^ string_of_jlite_type m.rettype 
					^ string_of_jlite_type t1 ^ ". \n")
				else true
			in { m with stmts=newstmts;
				}

(* TypeCheck a JLite Program. 
   Return a new JLite Program 
   where expressions are annotated with types
*)
let type_check_jlite_program  
	(p:jlite_program) : jlite_program=
	let type_check_class_main 
		((cname,mmthd):class_main ) =
		(cname,(type_check_mthd_decl p [] cname mmthd )) in
	let rec type_check_class_decl 
		((cname,cvars,cmthds):class_decl) =
		(* TypeCheck field declarations *)
		let (retval, errmsg) = 
			(type_check_var_decl_list p cvars) in
		if (retval==false) then 
			failwith 
			("\nType-check error in " ^ cname 
			^ " field declarations." ^ errmsg ^ "\n")
		(* TypeCheck methods overloading *)
		else let (retval, errmsg) = 
			(type_check_md_overloading cname cmthds) in
			if (retval==false) then 
				failwith 
				("\nType-check error in " ^ cname 
				^ " method declarations." ^ errmsg ^ "\n")
			(* TypeCheck method declarations *)
			else let env = (create_scoped_var_decls cvars 1) 
			in (cname, cvars, 
				List.map 
				(type_check_mthd_decl p env cname) cmthds
				)
	in 
	begin
		let (mainclass, classes) = p in 
		let newmain = 
			(type_check_class_main mainclass) in
		let newclasses=
			(List.map type_check_class_decl classes) in
		(newmain, newclasses)
	end
